import torch
from cs336_basics.tokenizer import Tokenizer
from cs336_basics.model import Transformer


def decode(
    model: Transformer,           # 用于生成文本的Transformer模型
    tokenizer: Tokenizer,         # 用于编码和解码token的分词器
    prompt: str,                  # 生成文本的起始提示
    max_new_tokens: int = 32,     # 最大生成token数
    temperature: float = 0.7,     # 控制生成随机性的温度参数
    top_p: float = 0.9,           # 用于nucleus sampling的累积概率阈值
):
    # 获取结束token的ID（空字符串的编码）
    end_id = tokenizer.encode("")[0]
    # 将提示文本编码为token ID序列
    input_ids = tokenizer.encode(prompt)
    # 获取模型所在的设备（CPU或GPU）
    device = next(model.parameters()).device
    # 获取模型的最大上下文长度
    context_length = model.context_length

    # 禁用梯度计算，因为这是推理过程
    with torch.no_grad():
        # 循环生成新token，直到达到最大生成数量
        for _ in range(max_new_tokens):
            # 如果输入序列超过模型上下文长度，则只取最后context_length个token
            window_input_ids = input_ids[-context_length:] if len(input_ids) >= context_length else input_ids
            # 将输入序列转换为torch张量并移动到指定设备
            x = torch.tensor([window_input_ids], dtype=torch.long, device=device)

            # 通过模型获取logits输出
            logits = model(x)
            # 取最后一个位置的logits作为下一个token的预测
            next_logits = logits[0, -1, :]

            # 应用温度缩放来调整分布的锐度
            scaled = next_logits / temperature
            # 为了数值稳定性，减去最大值
            stable = scaled - scaled.max()
            # 计算指数值
            exp_vals = stable.exp()
            # 计算概率分布
            probs = exp_vals / exp_vals.sum()

            # 对概率进行降序排序
            sorted_probs, sorted_idxs = torch.sort(probs, descending=True)
            # 计算累积概率
            cumsum = torch.cumsum(sorted_probs, dim=0)
            # 找到累积概率超过top_p的位置
            cutoff_idx = torch.searchsorted(cumsum, top_p)
            # 截取top-p部分的概率和对应的token索引
            trimmed_probs = sorted_probs[: cutoff_idx + 1]
            trimmed_idxs = sorted_idxs[: cutoff_idx + 1]
            # 重新归一化概率
            trimmed_probs /= trimmed_probs.sum()

            # 从截取的概率分布中采样下一个token
            next_token = trimmed_idxs[torch.multinomial(trimmed_probs, 1).item()]
            # 如果生成了结束token，则停止生成
            if next_token.item() == end_id:
                break
            # 将新生成的token添加到输入序列中
            input_ids.append(next_token.item())

    # 将完整的token序列解码为文本并返回
    return tokenizer.decode(input_ids)